package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"time"

	"adkr/pkg/turritopsis"
	"adkr/pkg/adkr"
)

func main() {
	var (
		epochs = flag.Int("epochs", 30, "Number of epochs to run")
		r      = flag.Int("r", 5, "ADKR execution interval (every R epochs)")
		n      = flag.Int("n", 4, "Number of nodes")
		f      = flag.Int("f", 1, "Number of faulty nodes")
		pid    = flag.Int("pid", 0, "Process ID")
		sid    = flag.String("sid", "turritopsis-test", "Session ID")
	)
	flag.Parse()

	logger := log.New(os.Stdout, fmt.Sprintf("[PID-%d] ", *pid), log.LstdFlags)
	logger.Printf("Starting Turritopsis test: epochs=%d, R=%d, N=%d, F=%d", *epochs, *r, *n, *f)

	// Create configuration with R-based ADKR execution
	cfg := turritopsis.Config{
		SID:    *sid,
		PID:    *pid,
		N:      *n,
		F:      *f,
		Epochs: *epochs,
		R:      *r, // ADKR execution interval
	}

	// Set up ACS function (stub for now, can be replaced with sDumbo)
	cfg.ACS = func(sid string, input []byte) ([][]byte, error) {
		logger.Printf("ACS[%s]: processing input", sid)
		// Simulate ACS processing time
		time.Sleep(10 * time.Millisecond)
		// Return the input as a batch
		if input != nil {
			return [][]byte{input}, nil
		}
		return [][]byte{}, nil
	}

	// Set up simplified Key Refresh function with R-based execution
	cfg.KeyRefreshOne = func(epoch int) (*adkr.KeyRefreshOutput, error) {
		// Only execute ADKR every R epochs
		if epoch%*r != 0 {
			logger.Printf("Epoch %d: skipping ADKR (not R=%d interval)", epoch, *r)
			return nil, nil
		}

		logger.Printf("Epoch %d: executing simplified ADKR key refresh", epoch)
		
		// Simulate ADKR processing time
		start := time.Now()
		time.Sleep(100 * time.Millisecond) // Simulate ADKR processing
		duration := time.Since(start)
		
		logger.Printf("Epoch %d: simplified ADKR completed in %v", epoch, duration)
		
		// Return a mock KeyRefreshOutput
		return &adkr.KeyRefreshOutput{
			ConfigChain: [][]int{{0, 1, 2, 3}}, // Same config, no change
		}, nil
	}

	// Set up proof chain builder
	cfg.AfterEpoch = func(epoch int, plan turritopsis.EpochPlan, krResult *adkr.KeyRefreshOutput) error {
		// Log epoch completion with proof information
		krExecuted := epoch%*r == 0
		logger.Printf("Epoch %d completed: N=%d, F=%d, L=%d, KR executed=%v", 
			epoch, plan.N, plan.F, plan.L, krExecuted)
		
		// Log proof information if KR was executed
		if krExecuted && krResult != nil {
			logger.Printf("Epoch %d proof: config_chain=%v", epoch, krResult.ConfigChain)
		}
		return nil
	}

	// Create and run protocol
	proto := turritopsis.New(cfg, logger)
	if err := proto.Run(); err != nil {
		logger.Fatalf("Protocol failed: %v", err)
	}

	logger.Printf("Turritopsis test completed successfully")
	
	// Print summary
	logger.Printf("Test completed: %d epochs, ADKR executed every %d epochs", *epochs, *r)
	logger.Printf("ADKR was executed %d times (epochs: 0, 5, 10, 15, 20, 25)", (*epochs+*r-1)/(*r))
}
